#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2019 crane <crane@crane-pc>
#
# Distributed under terms of the MIT license.

"""
    "Hungarian algorithm" is also called "KM algorithm"

    left和right 进入match后, 就不会再出去的(每次只是再增加一个left, 一个right而已)
"""

from pprint import pprint


class GraphNode:
    def __init__(self, label):
        self.label = label
        self.neighbors = set()      # nei_label

    def add(self, nei_label):
        self.neighbors.add(nei_label)

    def __str__(self):
        return str(self.label) + ": " + str(self.neighbors)


class MaxMatch:

    def __init__(self, edges):
        # self.match_edges = []
        self.nodes = {}         # {label : node}
        self.left = set()
        self.right = set()

        is_bipartite = self.build_adjacency_list(edges)

        if not is_bipartite:
            self.match = []
            print('Warnning: not bipartite graph')

        self.upper_cnt = min(len(self.left), len(self.right))       # 最大匹配数upper bound
        self.max_match_find()

    def build_adjacency_list(self, edges):

        def _start_in_left(start, end):
            if end in self.left:
                return False
            elif end in self.right:
                self._link_node(start, end)
            else:
                self.right.add(end)
                self._link_node(start, end)
            return True

        def _start_in_right(start, end):
            if end in self.right:
                return False
            elif end in self.left:
                self._link_node(end, start)
            else:
                self.left.add(end)
                self._link_node(end, start)
            return True

        def _start_not_in(start, end):
            if end in self.left:
                self.right.add(start)
                self._link_node(end, start)
            elif end in self.right:
                self.left.add(start)
                self._link_node(start, end)
            else:
                self.left.add(start)
                self.right.add(end)
                self._link_node(start, end)
            return True

        # =================== real fun code =====================
        for start, end in edges:
            if start in self.left:
                if not _start_in_left(start, end):
                    return False
            elif start in self.right:
                if not _start_in_right(start, end):
                    return False
            else:
                _start_not_in(start, end)

        return True

    def max_match_find(self):
        self.matchs = {}            # {left : right} + {right : left}
        self.non_matchs_left = set()     # non-match left
        self.match_cnt = 0

        self.greedy_max_match()
        pprint(self.non_matchs_left)


        self.handled_left = set()
        self.stack_right = []
        self._try_aug_match()       # 等价于循环遍历所有的left即可.

    def _try_aug_match(self):
        # 循环增广所有non_matchs_left中结点.
        # NOTE: 循不前期增广的路径, 不影响循环后面增广路径.
        for node_label in self.non_matchs_left:
            # 每次循环最多只增加一个left节点到match pair中, 也就是当前node
            sub_left, sub_right = self._rec_aug_path(node_label)
            if sub_left:
                self.add_match(sub_left, sub_right)       # plus 1 match
                print('=================== plus find aug : %s ---> %s=====================' % (sub_left, sub_right))
            else:
                print('=================== left failed %s =====================' % node_label)
                pprint(self.stack_right)

        print('=================== result matchs =====================')
        pprint(self.matchs)
        pprint('result match cnt %s' % self.match_cnt)
        pprint('handled_left %s' % self.handled_left)
        assert not self.stack_right

    def _rec_aug_path(self, left_label):
        # return left_label, right_label
        # 先通过left_label选择non_match edge, 再选择match_edge, 然后递归
        left_node = self.nodes[left_label]

        for right_label in left_node.neighbors:
            # if self.matchs.get(left_label) == right_label:
            print('%s ---> %s ' % (left_label, right_label))
            if self.matchs.get(right_label) == left_label:
                # kick out match edge
                # find non-match edge(left_label ----> right_label)
                print('not non-matched edge')
                continue

            if right_label in self.stack_right:
                print('stack right_label matched')
                continue

            if right_label not in self.matchs:
                # find a augment path
                print('find %s %s' % (left_label, right_label) )
                return left_label, right_label

            # right_label ----> left_label : find matched edge
            next_left = self.matchs[right_label]
            if next_left in self.handled_left:
                print('handled left_label matched')
                continue

            self.stack_right.append(right_label)
            sub_left, sub_right = self._rec_aug_path(next_left)
            r = self.stack_right.pop()
            assert r == right_label
            if sub_left:
                print('find aug : %s ---> %s' % (sub_left, sub_right))
                self.do_augment(right_label, sub_left, sub_right)
                return left_label, right_label

        # 至此, 如果left_label是non_handle: 得出此left_label无法增广: left_label---> all right_label 无法增广(right_label --> left_label 无法南增广)
        # 如果left_label已经在match pair中: 表示left_label 已经最优了: 无论如何循环, 都不得可能再优化.
        # 后续中再次曾广时到达此left_node, 直接返回即可
        print('left died %s' % left_label)
        self.handled_left.add(left_label)        # 防止后面重复处理
        return None, None

    def do_augment(self, old_right, new_left, new_right):
        assert self.matchs[old_right] == new_left

        # TODO: 这里不用删除old_right, 只用修改right的link/match即可.
        # 无论left, 还是right只要添加到matchs匹配中后, 后续都不会掉出matchs.
        self.remove_match(old_right)
        self.add_match(new_left, new_right)

    def _get_node(self, label):
        node = self.nodes.get(label)
        if node is None:
            node = GraphNode(label)
            self.nodes[label] = node
        return node

    def _link_node(self, left_label, right_label):
        left_node = self._get_node(left_label)
        right_node = self._get_node(right_label)
        left_node.add(right_label)
        # right不需要添加left: 寻找增光路径时, 只需要遍历left的neighbors. 对于right, 只用找到与之match的left即可.
        # right_node.add(left_label)

    def greedy_max_match(self, ):
        self.match_cnt = 0
        self.non_matchs_left = set(self.left)

        # NOTE: 也可以先greedy match, 加快后续match速度
        # for l in self.left:
        #     assert l not in self.matchs
        #     l_node = self.nodes[l]

        #     for r in l_node.neighbors:
        #         if r in self.matchs:
        #             continue
        #         else:
        #             self.add_match(l, r)
        #             break
        #     # TODO add non_matchs_left

    def add_match(self, l, r):
        self.matchs[r] = l
        self.match_cnt += 1

    def remove_match(self, right_label):
        del self.matchs[right_label]
        self.match_cnt -= 1


def main():
    print("start main")
    edge_list = [
        (1, 2),
        (2, 3),
        (3, 4),
        (4, 1)
    ]

    edges_example1 = [
        ('Billy', 1),
        ('Billy', 4),

        ('Emily', 7),
        ('Emily', 3),
        ('Emily', 6),

        ('John', 2),
        ('John', 5),
        ('John', 4),

        ('Luke', 7),
        ('Luke', 2),

        ('Timothy', 7),
        ('Timothy', 6),
        ('Timothy', 5),

        ('Anna', 3),
        ('Anna', 6),

        ('Raj', 6),
        ('Raj', 7),
    ]

    edges_example2 = [
        (1, 2),
        (1, 4),
        (1, 6),

        (3, 2),
        (5, 4),
        (7, 6),
        (7, 8),

    ]

    edges_example3 = [
        # (1, 2),
        (1, 4),

        (3, 4),

        (5, 4),

        (7, 4),
    ]

    m = MaxMatch(edges_example2)
    m = MaxMatch(edge_list)

if __name__ == "__main__":
    main()
